package ace.handling;

import gov.nist.atlas.ATLASAccessException;
import gov.nist.atlas.ATLASClass;
import gov.nist.atlas.ATLASElement;
import gov.nist.atlas.Analysis;
import gov.nist.atlas.Anchor;
import gov.nist.atlas.Annotation;
import gov.nist.atlas.Children;
import gov.nist.atlas.ConstraintSet;
import gov.nist.atlas.CorporaManager;
import gov.nist.atlas.Corpus;
import gov.nist.atlas.Id;
import gov.nist.atlas.IdentifiableATLASElement;
import gov.nist.atlas.Metadata;
import gov.nist.atlas.Region;
import gov.nist.atlas.ReusableATLASElement;
import gov.nist.atlas.Signal;
import gov.nist.atlas.ValidationResult;
import gov.nist.atlas.event.ContentChangeListener;
import gov.nist.atlas.event.ContentWillChangeListener;
import gov.nist.atlas.impl.AnalysisImpl;
import gov.nist.atlas.impl.AnnotationImpl;
import gov.nist.atlas.impl.CorpusImpl;
import gov.nist.atlas.io.ATLASExport;
import gov.nist.atlas.io.ATLASImport;
import gov.nist.atlas.io.xml.AIFXMLExport;
import gov.nist.atlas.io.xml.AIFXMLImport;
import gov.nist.atlas.spi.TypeImplementationDelegate;
import gov.nist.atlas.type.ATLASType;
import gov.nist.atlas.type.ChildrenType;
import gov.nist.atlas.type.CorpusType;
import gov.nist.atlas.type.OptionalChildrenDefinition;
import gov.nist.atlas.type.RequiredChildDefinition;
import gov.nist.atlas.util.ATLASElementSet;
import gov.nist.atlas.util.ATLASImplementation;
import gov.nist.atlas.util.Mutable;
import gov.nist.atlas.util.MutableATLASElementSet;
import gov.nist.maia.MAIAScheme;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.StringReader;
import java.io.Writer;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

import org.mitre.jawb.atlas.AWBATLASImplementation;
import org.mitre.jawb.atlas.AWBAnnotationImpl;
import org.mitre.jawb.atlas.AWBCorpusImpl;
import org.mitre.jawb.atlas.AWBDocument;
import org.mitre.jawb.atlas.JAWB;
import org.mitre.jawb.atlas.MultiPhraseAnnotation;
import org.mitre.jawb.atlas.SubordinateSetsAnnotation;
import org.mitre.jawb.io.SgmlDocument;
import org.mitre.jawb.io.SgmlElement;
import org.mitre.jawb.tasks.DefaultInlineExporter;
import org.mitre.jawb.tasks.DefaultInlineImporter;
import org.mitre.jawb.tasks.Exporter;
import org.mitre.jawb.tasks.Importer;
import org.mitre.jawb.tasks.Task;
import org.mitre.jawb.tasks.generic.GenericTask;

public class ACEHandler {

	private static Set<Id> annotatedEntities = new HashSet<Id>();
	private static Set<String> annotatedRelations = new HashSet<String>();
	private static Set<String> relationsEntityCombinations = new HashSet<String>();


	/**
	 * @param args
	 * @throws IOException 
	 * @throws URISyntaxException 
	 */
	public static void main(String[] args) throws IOException, URISyntaxException {
		
		String sgmDatasetsFolder = "/home/pjbarrio/Dataset/SGML-ACE-COREF/";
		
		String finalDirectory = "/0/taggedtestingDirectory/";
		
		URI file = new URI("file:///home/pjbarrio/Software/Callisto/data/ACE/VOA19980414.1700.3225.sgm.tmx.rdc.apf.xml.txt.aif.xml");
		
		List<Task> t = AWBDocument.getSupportingTasks(file);
		
		Task rdcTask = t.get(0);
		
		String fileList = "/home/pjbarrio/Dataset/ACE2003/ace_phase2/data/ace2_devtest/toImport/corpus_list_aif";

		String directory = "/home/pjbarrio/Dataset/ACE2003/ace_phase2/data/ace2_devtest/toImport/";
		
		File corpusAIF = new File(fileList);
		
		BufferedReader br = new BufferedReader(new FileReader(corpusAIF));
		
		String line = br.readLine();
		
		while (line!=null){
		
			System.err.println(line);
			
			URI fileToOpen = new URI("file://" + directory + line);
						
			AWBDocument document = AWBDocument.fromAIF(fileToOpen, rdcTask);
			
			Document doc = new Document(line.replace(".sgm.aif.xml", ""),document.getSignal().getCharsAt(0));
			
			Iterator f = document.getAllAnalyses();
			
			while (f.hasNext()){
				
				AnalysisImpl aux = (AnalysisImpl)f.next();
				
//				System.out.println("TYPE NAME: " + aux.getAnalysisType().getName());
				
				ATLASElementSet annotations = aux.getAllAnnotations();
				
				Iterator annotationsIterator = annotations.iterator();
				
				while (annotationsIterator.hasNext()){
					
					if (aux.getAnalysisType().getName().equals("ace_mention-set")){
						
						MultiPhraseAnnotation mpa = (MultiPhraseAnnotation)annotationsIterator.next();
	//					
	//					System.out.println("MENTION-SET: " + mpa.asAIFString());
	//					
	//					Region full = mpa.getRegion().getRegionWithRole("full");
	//					
	//					Region head = mpa.getRegion().getRegionWithRole("head");
	//					
	//					System.out.println("REGION: " + mpa.getRegion().asAIFString());
	//					
	//					System.out.println("FULL: " + full.asAIFString());
	//					
	//					System.out.println("HEAD: " + head.asAIFString());
						
					}else{
						
						//ace_rdc-mention-relation-set
						
						SubordinateSetsAnnotation ssa = (SubordinateSetsAnnotation)annotationsIterator.next();
						
//						System.out.println("OTRO: " + ssa.asAIFString());
											
//						System.out.println("REGION: " + ssa.getRegion().asAIFString());
						
						if (ssa.getAnnotationType().getName().equals("ace_edt-entity")){
						
							Entity entity = doc.getEntityFactory().generateEntity(ssa);
												
//							System.out.println("TYPE: " + ssa.getContent().getParameterWithRole("ace_edt-type").getValueAsString());
							
							Annotation primaryMention = ssa.getRegion().getAnnotationWithRole("ace_primary-mention");
							
//							System.out.println("PRIMARY MENTION: " + primaryMention.asAIFString());
							
							ATLASElementSet mentions = ssa.getRegion().getSubordinateSet(primaryMention.getATLASType());
							
							Iterator mentionsIterator = mentions.iterator();
							
							while (mentionsIterator.hasNext()){
								
								Annotation mention = (Annotation)mentionsIterator.next();
								
//								System.out.println("MENTIONS: " + mention.asAIFString());
								
								Entity mentionedEntity = doc.getEntityFactory().generateEntity(mention);
								
								entity.associate(mentionedEntity);
								
//								System.out.println("REGION: " + mention.getRegion().asAIFString());
								
//								Region full = mention.getRegion().getRegionWithRole("full");
//								
//								Anchor fullStart = full.getAnchorWithRole("start");
//								
//								Anchor fullEnd = full.getAnchorWithRole("end");
//								
//								Region head = mention.getRegion().getRegionWithRole("head");
//								
//								Anchor headStart = head.getAnchorWithRole("start");
//								
//								Anchor headEnd = head.getAnchorWithRole("end");
								
//								System.out.println("FULL: " + full.asAIFString());
								
//								System.out.println(fullStart.asAIFString());
//								
//								System.out.println(fullEnd.asAIFString());
//								
//								System.out.println("HEAD: " + head.asAIFString());
//								
//								System.out.println(headStart.asAIFString());
//								
//								System.out.println(headEnd.asAIFString());
							}
							
							
						} else {
							
							if (ssa.getAnnotationType().getName().equals("ace_rdc-mention-relation")){
								
								Relation relation = doc.getRelationFactory().generateRelation(ssa);
								
//								System.out.println("TYPE: " + ssa.getContent().getParameterWithRole("ace_rdc-type").getValueAsString());
								
								Annotation argument1 = ssa.getRegion().getAnnotationWithRole("ace_rdc-argument1");
								
								Entity arg1 = doc.getEntityFactory().generateEntity(argument1);
								
								relation.setArgument1(arg1);
								
//								System.out.println("ARGUMENT1: " + argument1.asAIFString());
								
//								System.out.println("REGION: " + argument1.getRegion().asAIFString());
								
								Annotation argument2 = ssa.getRegion().getAnnotationWithRole("ace_rdc-argument2");
								
								Entity arg2 = doc.getEntityFactory().generateEntity(argument2);
								
								relation.setArgument2(arg2);
								
//								System.out.println("ARGUMENT2: " + argument2.asAIFString());
//								
//								System.out.println("REGION: " + argument2.getRegion().asAIFString());
								
							}
							
						}
						
					}
						
				}
							
			}	
			
			saveSGMLDocumentInDatasets(doc,sgmDatasetsFolder,finalDirectory);
			
			System.gc();
			
			line = br.readLine();
			
		}
				
	}

	private static void saveSGMLDocumentInDatasets(Document doc,
			String sgmDatasetsFolder, String finalDirectory) throws IOException {
		
		String[] relations = doc.getRelationsList();
		
		for (String rel : relations) {
		
			clearAnnotations();
			
			String directory = sgmDatasetsFolder + rel + "/" + finalDirectory;
			
			SgmlDocument sgmlDocument = new SgmlDocument(new StringReader(doc.getText()));
			
			addEntities(sgmlDocument,doc.getEntities());
			
			new File(directory).mkdirs();
			
			int index = 1;
			
			List<Relation> rels = doc.getRelations(rel);
			
			for (Relation relation : rels) {
				
				annotateArgument(relation.getArgument1(),sgmlDocument);
				
				annotateArgument(relation.getArgument2(),sgmlDocument);
				
				annotateRelation(relation,sgmlDocument,index,rel);
				
				String combination = rel + " - " + relation.getArgument1().getTag() + " - " + relation.getArgument2().getTag();
				
				if (!relationsEntityCombinations.contains(combination)){
					
					System.out.println(combination);
					
					relationsEntityCombinations.add(combination);
					
				}
				
				index++;
				
			}
			
			saveDocument(new String(directory + doc.getName() + ".sgm"), sgmlDocument);
			
		}
		
		
	}

	private static void addEntities(SgmlDocument sgmlDocument,
			Enumeration<Entity> entities) {
		for (; entities.hasMoreElements(); ) {
			
			Entity entity = entities.nextElement();
			
			annotateArgument(entity, sgmlDocument);
			
		}
	}

	private static void clearAnnotations() {
		
		annotatedEntities.clear();
		
		annotatedRelations.clear();
		
	}

	private static void saveDocument(String file, SgmlDocument sgmlDocument) throws IOException {
		
		sgmlDocument.writeSgml(new BufferedWriter(new FileWriter(new File(file))));
		
	}

	private static void annotateRelation(Relation relation,
			SgmlDocument sgmlDocument, int index, String tag) {
		
		if (!alreadyAnnotated(relation)){
		
			annotate(relation);
			
			Entity a1 = relation.getArgument1();
			
			SgmlElement arg1 = sgmlDocument.createContentTag(a1.getStart(), a1.getEnd(), tag + "First", true);
			
			arg1.putAttribute("id", Integer.toString(index));
			
//			sgmlDocument.addElement(arg1, false);
			
			Entity a2 = relation.getArgument2();
			
			SgmlElement arg2 = sgmlDocument.createContentTag(a2.getStart(), a2.getEnd(), tag + "Second", true);
			
			arg2.putAttribute("id", Integer.toString(index));
			
//			sgmlDocument.addElement(arg2, false);		

		}
	}

	private static void annotate(Relation relation) {
		
		annotatedRelations.add(relation.getId());
		
	}

	private static boolean alreadyAnnotated(Relation relation) {
		return annotatedRelations.contains(relation.getId());
	}

	private static void annotateArgument(Entity entity,
			SgmlDocument sgmlDocument) {
		
		if (!alreadyAnnotated(entity)){
		
			annotate(entity);
			
			SgmlElement element = sgmlDocument.createContentTag(entity.getStart(), entity.getEnd(), entity.getTag().toUpperCase(), false);
			
//			sgmlDocument.addElement(element, false);
			
		}
	
	}

	private static boolean alreadyAnnotated(Entity entity) {
		return annotatedEntities.contains(entity.getId());
	}

	private static void annotate(Entity entity) {
		
		annotatedEntities.add(entity.getId());
		
	}

}
